#include <linux/module.h>
#include <linux/platform_device.h>
#include <linux/miscdevice.h>
#include <linux/fs.h>
#include <linux/uaccess.h>
#include <linux/slab.h>
#include <linux/mutex.h>
#include <linux/of.h>

#include "twclml_drv.h"

/* 模拟的设备列表，包含4个设备 */
static struct twcl_nvmlDevice g_devices[] = {
    {
        .index = 0,
        .pciInfo = {
            .busId = "0000:01:00.0",
            .domain = 0x0000,
            .bus = 0x01,
            .device = 0x00,
            .function = 0x0
        },
        .computeMajor = 8,
        .computeMinor = 6,
        .numNvLinks = 4,
        .chipsetSupportsP2P = true,
        .gpuSupportsP2P = true,
        .iohTopologySupportsP2P = true,
        .p2pDisabledByRegKey = false,
        .remotePcie = {
            .numLinks = 2,
            .remotePci = {
                { .busId = "0000:02:00.0", .domain = 0x0000, .bus = 0x02, .device = 0x00, .function = 0x0 },
                { .busId = "0000:03:00.0", .domain = 0x0000, .bus = 0x03, .device = 0x00, .function = 0x0 },
                { .busId = "", .domain = 0, .bus = 0, .device = 0, .function = 0 },
                { .busId = "", .domain = 0, .bus = 0, .device = 0, .function = 0 }
            },
            .linkState = { NVML_ENABLED, NVML_ENABLED, NVML_DISABLED, NVML_DISABLED },
            .linkCaps = {
                { .p2pSupported = 1, .sysmemAccess = 1, .p2pAtomics = 1, .sysmemAtomics = 0, .sliBridge = 0, .valid = 1 },
                { .p2pSupported = 1, .sysmemAccess = 1, .p2pAtomics = 0, .sysmemAtomics = 0, .sliBridge = 1, .valid = 1 },
                {0}, {0}
            },
            .p2pCaps = {
                { .read = NVML_P2P_STATUS_OK, .write = NVML_P2P_STATUS_OK, .nvlink = NVML_P2P_STATUS_OK, .atomics = NVML_P2P_STATUS_OK, .prop = NVML_P2P_STATUS_OK },
                { .read = NVML_P2P_STATUS_OK, .write = NVML_P2P_STATUS_OK, .nvlink = NVML_P2P_STATUS_OK, .atomics = NVML_P2P_STATUS_NOT_SUPPORTED, .prop = NVML_P2P_STATUS_OK },
                {0}, {0}
            }
        }
    },
    {
        .index = 1,
        .pciInfo = {
            .busId = "0000:02:00.0",
            .domain = 0x0000,
            .bus = 0x02,
            .device = 0x00,
            .function = 0x0
        },
        .computeMajor = 8,
        .computeMinor = 6,
        .numNvLinks = 4,
        .chipsetSupportsP2P = true,
        .gpuSupportsP2P = true,
        .iohTopologySupportsP2P = true,
        .p2pDisabledByRegKey = false,
        .remotePcie = {
            .numLinks = 2,
            .remotePci = {
                { .busId = "0000:01:00.0", .domain = 0x0000, .bus = 0x01, .device = 0x00, .function = 0x0 },
                { .busId = "0000:04:00.0", .domain = 0x0000, .bus = 0x04, .device = 0x00, .function = 0x0 },
                { .busId = "", .domain = 0, .bus = 0, .device = 0, .function = 0 },
                { .busId = "", .domain = 0, .bus = 0, .device = 0, .function = 0 }
            },
            .linkState = { NVML_ENABLED, NVML_ENABLED, NVML_DISABLED, NVML_DISABLED },
            .linkCaps = {
                { .p2pSupported = 1, .sysmemAccess = 1, .p2pAtomics = 1, .sysmemAtomics = 0, .sliBridge = 0, .valid = 1 },
                { .p2pSupported = 1, .sysmemAccess = 1, .p2pAtomics = 0, .sysmemAtomics = 0, .sliBridge = 1, .valid = 1 },
                {0}, {0}
            },
            .p2pCaps = {
                { .read = NVML_P2P_STATUS_OK, .write = NVML_P2P_STATUS_OK, .nvlink = NVML_P2P_STATUS_OK, .atomics = NVML_P2P_STATUS_OK, .prop = NVML_P2P_STATUS_OK },
                { .read = NVML_P2P_STATUS_OK, .write = NVML_P2P_STATUS_OK, .nvlink = NVML_P2P_STATUS_OK, .atomics = NVML_P2P_STATUS_NOT_SUPPORTED, .prop = NVML_P2P_STATUS_OK },
                {0}, {0}
            }
        }
    },
    {
        .index = 2,
        .pciInfo = {
            .busId = "0000:03:00.0",
            .domain = 0x0000,
            .bus = 0x03,
            .device = 0x00,
            .function = 0x0
        },
        .computeMajor = 7,
        .computeMinor = 5,
        .numNvLinks = 2,
        .chipsetSupportsP2P = false,
        .gpuSupportsP2P = true,
        .iohTopologySupportsP2P = true,
        .p2pDisabledByRegKey = false,
        .remotePcie = {
            .numLinks = 2,
            .remotePci = {
                { .busId = "0000:01:00.0", .domain = 0x0000, .bus = 0x01, .device = 0x00, .function = 0x0 },
                { .busId = "0000:04:00.0", .domain = 0x0000, .bus = 0x04, .device = 0x00, .function = 0x0 },
                { .busId = "", .domain = 0, .bus = 0, .device = 0, .function = 0 },
                { .busId = "", .domain = 0, .bus = 0, .device = 0, .function = 0 }
            },
            .linkState = { NVML_ENABLED, NVML_ENABLED, NVML_DISABLED, NVML_DISABLED },
            .linkCaps = {
                { .p2pSupported = 1, .sysmemAccess = 1, .p2pAtomics = 0, .sysmemAtomics = 0, .sliBridge = 0, .valid = 1 },
                { .p2pSupported = 1, .sysmemAccess = 1, .p2pAtomics = 0, .sysmemAtomics = 0, .sliBridge = 1, .valid = 1 }
            },
            .p2pCaps = {
                { .read = NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED, .write = NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED, .nvlink = NVML_P2P_STATUS_OK, .atomics = NVML_P2P_STATUS_NOT_SUPPORTED, .prop = NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED },
                { .read = NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED, .write = NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED, .nvlink = NVML_P2P_STATUS_OK, .atomics = NVML_P2P_STATUS_NOT_SUPPORTED, .prop = NVML_P2P_STATUS_CHIPSET_NOT_SUPPORED }
            }
        }
    },
    {
        .index = 3,
        .pciInfo = {
            .busId = "0000:04:00.0",
            .domain = 0x0000,
            .bus = 0x04,
            .device = 0x00,
            .function = 0x0
        },
        .computeMajor = 8,
        .computeMinor = 0,
        .numNvLinks = 4,
        .chipsetSupportsP2P = true,
        .gpuSupportsP2P = true,
        .iohTopologySupportsP2P = true,
        .p2pDisabledByRegKey = false,
        .remotePcie = {
            .numLinks = 2,
            .remotePci = {
                { .busId = "0000:02:00.0", .domain = 0x0000, .bus = 0x02, .device = 0x00, .function = 0x0 },
                { .busId = "0000:03:00.0", .domain = 0x0000, .bus = 0x03, .device = 0x00, .function = 0x0 },
                { .busId = "", .domain = 0, .bus = 0, .device = 0, .function = 0 },
                { .busId = "", .domain = 0, .bus = 0, .device = 0, .function = 0 }
            },
            .linkState = { NVML_ENABLED, NVML_ENABLED, NVML_DISABLED, NVML_DISABLED },
            .linkCaps = {
                { .p2pSupported = 1, .sysmemAccess = 1, .p2pAtomics = 1, .sysmemAtomics = 1, .sliBridge = 0, .valid = 1 },
                { .p2pSupported = 1, .sysmemAccess = 1, .p2pAtomics = 1, .sysmemAtomics = 0, .sliBridge = 1, .valid = 1 },
                {0}, {0}
            },
            .p2pCaps = {
                { .read = NVML_P2P_STATUS_OK, .write = NVML_P2P_STATUS_OK, .nvlink = NVML_P2P_STATUS_OK, .atomics = NVML_P2P_STATUS_OK, .prop = NVML_P2P_STATUS_OK },
                { .read = NVML_P2P_STATUS_OK, .write = NVML_P2P_STATUS_OK, .nvlink = NVML_P2P_STATUS_OK, .atomics = NVML_P2P_STATUS_OK, .prop = NVML_P2P_STATUS_OK },
                {0}, {0}
            }
        }
    }
};

#define NUM_DEVICES (sizeof(g_devices) / sizeof(g_devices[0]))

/* NVML库函数实现 */
nvmlReturn_t nvmlInit(void) {
    return NVML_SUCCESS;
}

nvmlReturn_t nvmlInit_v2(void) {
    return NVML_SUCCESS;
}

nvmlReturn_t nvmlShutdown(void) {
    return NVML_SUCCESS;
}

nvmlReturn_t nvmlDeviceGetCount(unsigned int* count) {
    if (!count) return NVML_ERROR_INVALID_ARGUMENT;
    *count = NUM_DEVICES;
    return NVML_SUCCESS;
}

nvmlReturn_t nvmlDeviceGetCount_v2(unsigned int* count) {
    return nvmlDeviceGetCount(count);
}

nvmlReturn_t nvmlDeviceGetHandleByIndex(unsigned int index, nvmlDevice_t* device) {
    if (!device || index >= NUM_DEVICES) return NVML_ERROR_INVALID_ARGUMENT;
    *device = &g_devices[index];
    return NVML_SUCCESS;
}

nvmlReturn_t nvmlDeviceGetPciInfo_v2(nvmlDevice_t device, nvmlPciInfo_t* pci) {
    if (!device || !pci) return NVML_ERROR_INVALID_ARGUMENT;
    *pci = device->pciInfo;
    return NVML_SUCCESS;
}

nvmlReturn_t nvmlDeviceGetCudaComputeCapability(nvmlDevice_t device, int* major, int* minor) {
    if (!device || !major || !minor) return NVML_ERROR_INVALID_ARGUMENT;
    *major = device->computeMajor;
    *minor = device->computeMinor;
    return NVML_SUCCESS;
}

nvmlReturn_t nvmlDeviceGetNvLinkState(nvmlDevice_t device, unsigned int link, 
                                     nvmlEnableState_t* isActive) {
    if (!device || !isActive || link >= device->numNvLinks) 
        return NVML_ERROR_INVALID_ARGUMENT;
    
    if (link < device->remotePcie.numLinks)
        *isActive = device->remotePcie.linkState[link];
    else
        *isActive = NVML_DISABLED;
    
    return NVML_SUCCESS;
}

nvmlReturn_t nvmlDeviceGetNvLinkRemotePciInfo(nvmlDevice_t device, unsigned int link, 
                                             nvmlPciInfo_t* pci) {
    if (!device || !pci || link >= device->numNvLinks)
        return NVML_ERROR_INVALID_ARGUMENT;
    
    if (link < device->remotePcie.numLinks)
        *pci = device->remotePcie.remotePci[link];
    else
        *pci = device->pciInfo;
    
    return NVML_SUCCESS;
}

static int findLinkBetweenDevices(nvmlDevice_t device1, nvmlDevice_t device2) {
    for (unsigned int i = 0; i < device1->remotePcie.numLinks; i++) {
        if (strcmp(device1->remotePcie.remotePci[i].busId, device2->pciInfo.busId) == 0) {
            return i;
        }
    }
    return -1;
}

nvmlReturn_t nvmlDeviceGetP2PStatus(nvmlDevice_t device1, nvmlDevice_t device2, 
                                   nvmlGpuP2PCapsIndex_t p2pIndex, 
                                   nvmlGpuP2PStatus_t* p2pStatus) {
    if (!device1 || !device2 || !p2pStatus || p2pIndex >= NVML_P2P_CAPS_INDEX_UNKNOWN)
        return NVML_ERROR_INVALID_ARGUMENT;
    
    if (device1 == device2) {
        *p2pStatus = NVML_P2P_STATUS_NOT_SUPPORTED;
        return NVML_SUCCESS;
    }
    
    int linkIdx = findLinkBetweenDevices(device1, device2);
    if (linkIdx == -1) {
        *p2pStatus = NVML_P2P_STATUS_NOT_SUPPORTED;
        return NVML_SUCCESS;
    }
    
    switch (p2pIndex) {
        case NVML_P2P_CAPS_INDEX_READ:
            *p2pStatus = device1->remotePcie.p2pCaps[linkIdx].read;
            break;
        case NVML_P2P_CAPS_INDEX_WRITE:
            *p2pStatus = device1->remotePcie.p2pCaps[linkIdx].write;
            break;
        case NVML_P2P_CAPS_INDEX_NVLINK:
            *p2pStatus = device1->remotePcie.p2pCaps[linkIdx].nvlink;
            break;
        case NVML_P2P_CAPS_INDEX_ATOMICS:
            *p2pStatus = device1->remotePcie.p2pCaps[linkIdx].atomics;
            break;
        case NVML_P2P_CAPS_INDEX_PROP:
            *p2pStatus = device1->remotePcie.p2pCaps[linkIdx].prop;
            break;
        default:
            *p2pStatus = NVML_P2P_STATUS_UNKNOWN;
            return NVML_ERROR_NOT_SUPPORTED;
    }
    
    return NVML_SUCCESS;
}

const char* nvmlErrorString(nvmlReturn_t r) {
    switch (r) {
        case NVML_SUCCESS: return "Success";
        case NVML_ERROR_UNINITIALIZED: return "Library not initialized";
        case NVML_ERROR_INVALID_ARGUMENT: return "Invalid argument";
        case NVML_ERROR_NOT_SUPPORTED: return "Not supported";
        case NVML_ERROR_NOT_FOUND: return "Not found";
        default: return "Unknown error";
    }
}

nvmlReturn_t nvmlDeviceGetHandleByPciBusId(const char* pciBusId, nvmlDevice_t* device) {
    if (!pciBusId || !device) return NVML_ERROR_INVALID_ARGUMENT;
    
    for (unsigned int i = 0; i < NUM_DEVICES; i++) {
        if (strcmp(pciBusId, g_devices[i].pciInfo.busId) == 0) {
            *device = &g_devices[i];
            return NVML_SUCCESS;
        }
    }
    
    return NVML_ERROR_NOT_FOUND;
}

nvmlReturn_t nvmlDeviceGetIndex(nvmlDevice_t device, unsigned int* index) {
    if (!device || !index) return NVML_ERROR_INVALID_ARGUMENT;
    
    for (unsigned int i = 0; i < NUM_DEVICES; i++) {
        if (&g_devices[i] == device) {
            *index = i;
            return NVML_SUCCESS;
        }
    }
    
    return NVML_ERROR_INVALID_ARGUMENT;
}

nvmlReturn_t nvmlDeviceGetNvLinkCapability(nvmlDevice_t device, unsigned int link, 
                                          nvmlNvLinkCapability_t capability, 
                                          unsigned int* capResult) {
    if (!device || !capResult || link >= device->numNvLinks || capability >= NVML_NVLINK_CAP_COUNT)
        return NVML_ERROR_INVALID_ARGUMENT;
    
    if (link >= device->remotePcie.numLinks)
        return NVML_ERROR_NOT_SUPPORTED;
    
    switch (capability) {
        case NVML_NVLINK_CAP_P2P_SUPPORTED:
            *capResult = device->remotePcie.linkCaps[link].p2pSupported;
            break;
        case NVML_NVLINK_CAP_SYSMEM_ACCESS:
            *capResult = device->remotePcie.linkCaps[link].sysmemAccess;
            break;
        case NVML_NVLINK_CAP_P2P_ATOMICS:
            *capResult = device->remotePcie.linkCaps[link].p2pAtomics;
            break;
        case NVML_NVLINK_CAP_SYSMEM_ATOMICS:
            *capResult = device->remotePcie.linkCaps[link].sysmemAtomics;
            break;
        case NVML_NVLINK_CAP_SLI_BRIDGE:
            *capResult = device->remotePcie.linkCaps[link].sliBridge;
            break;
        case NVML_NVLINK_CAP_VALID:
            *capResult = device->remotePcie.linkCaps[link].valid;
            break;
        default:
            return NVML_ERROR_NOT_SUPPORTED;
    }
    
    return NVML_SUCCESS;
}

/* 设备私有数据结构体 */
struct twclml_driver_data {
    struct miscdevice misc_dev;      /* misc设备结构体 */
    struct mutex lock;               /* 互斥锁 */
    int initialized;                 /* NVML库初始化状态 */
};

/* 全局驱动数据 */
static struct twclml_driver_data *twclml_drv_data;

/**
 * @brief 打开设备
 */
static int twclml_open(struct inode *inode, struct file *file)
{
    struct twclml_driver_data *data = twclml_drv_data;
    int ret = 0;
    
    /* 加锁保护 */
    mutex_lock(&data->lock);
    
    /* 如果未初始化，初始化NVML库 */
    if (!data->initialized) {
        ret = nvmlInit();
        if (ret != NVML_SUCCESS) {
            dev_err(data->misc_dev.this_device,  /* 修复：移除& */
                   "Failed to initialize TWCLML library: %s\n", 
                   nvmlErrorString(ret));
            ret = -EIO;
            goto out_unlock;
        }
        data->initialized = 1;
        dev_info(data->misc_dev.this_device, "TWCLML library initialized\n");  /* 修复：移除& */
    }
    
out_unlock:
    mutex_unlock(&data->lock);
    return ret;
}

/**
 * @brief 关闭设备
 */
static int twclml_release(struct inode *inode, struct file *file)
{
    struct twclml_driver_data *data = twclml_drv_data;
    
    /* 加锁保护 */
    mutex_lock(&data->lock);
    
    /* 关闭NVML库 */
    if (data->initialized) {
        nvmlShutdown();
        data->initialized = 0;
        dev_info(data->misc_dev.this_device, "TWCLML library shutdown\n");  /* 修复：移除& */
    }
    
    mutex_unlock(&data->lock);
    return 0;
}

/**
 * @brief IOCTL命令处理
 */
static long twclml_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
{
    struct twclml_driver_data *data = twclml_drv_data;
    int ret = 0;
    
    /* 检查命令合法性 */
    if (_IOC_TYPE(cmd) != TWCLML_IOCTL_MAGIC)
        return -ENOTTY;
    if (_IOC_NR(cmd) > TWCLML_IOCTL_MAXNR)
        return -ENOTTY;
    
    /* 加锁保护 */
    mutex_lock(&data->lock);
    
    /* 检查TWCLML库是否已初始化 */
    if (!data->initialized) {
        ret = -EINVAL;
        goto out_unlock;
    }
    
    /* 根据不同命令进行处理 */
    switch (cmd) {
        case TWCLML_IOCTL_GET_DEVICE_COUNT: {
            unsigned int count;
            ret = nvmlDeviceGetCount(&count);
            if (ret != NVML_SUCCESS) {
                dev_err(data->misc_dev.this_device,  /* 修复：移除& */
                       "Failed to get device count: %s\n", 
                       nvmlErrorString(ret));
                ret = -EIO;
                break;
            }
            
            /* 将结果拷贝到用户空间 */
            if (copy_to_user((void __user *)arg, &count, sizeof(count))) {
                ret = -EFAULT;
                break;
            }
            ret = 0;
            break;
        }
        
        case TWCLML_IOCTL_GET_DEVICE_INFO: {
            twclml_device_info_t info;
            nvmlDevice_t device;
            
            /* 从用户空间获取输入参数 */
            if (copy_from_user(&info, (void __user *)arg, sizeof(info))) {
                ret = -EFAULT;
                break;
            }
            
            /* 获取设备句柄 */
            ret = nvmlDeviceGetHandleByIndex(info.index, &device);
            if (ret != NVML_SUCCESS) {
                dev_err(data->misc_dev.this_device,  /* 修复：移除& */
                       "Failed to get device handle: %s\n", 
                       nvmlErrorString(ret));
                ret = -EIO;
                break;
            }
            
            /* 获取PCI信息 */
            {
                nvmlPciInfo_t pci;
                ret = nvmlDeviceGetPciInfo_v2(device, &pci);
                if (ret == NVML_SUCCESS) {
                    strncpy(info.bus_id, pci.busId, sizeof(info.bus_id)-1);
                    info.domain = pci.domain;
                    info.bus = pci.bus;
                    info.device = pci.device;
                    info.function = pci.function;
                }
            }
            
            /* 获取计算能力 */
            ret = nvmlDeviceGetCudaComputeCapability(device, 
                                                   &info.compute_major, 
                                                   &info.compute_minor);
            
            /* 获取NVLink数量 */
            info.num_nvlinks = device->numNvLinks;
            
            /* 将结果返回给用户空间 */
            if (copy_to_user((void __user *)arg, &info, sizeof(info))) {
                ret = -EFAULT;
                break;
            }
            ret = 0;
            break;
        }
        
        case TWCLML_IOCTL_GET_NVLINK_INFO: {
            twclml_nvlink_info_t info;
            nvmlDevice_t device;
            nvmlEnableState_t state;
            
            /* 从用户空间获取输入参数 */
            if (copy_from_user(&info, (void __user *)arg, sizeof(info))) {
                ret = -EFAULT;
                break;
            }
            
            /* 获取设备句柄 */
            ret = nvmlDeviceGetHandleByIndex(info.device_index, &device);
            if (ret != NVML_SUCCESS) {
                dev_err(data->misc_dev.this_device,  /* 修复：移除& */
                       "Failed to get device handle: %s\n", 
                       nvmlErrorString(ret));
                ret = -EIO;
                break;
            }
            
            /* 检查链路号是否有效 */
            if (info.link >= device->numNvLinks) {
                ret = -EINVAL;
                break;
            }
            
            /* 获取NVLink状态 */
            ret = nvmlDeviceGetNvLinkState(device, info.link, &state);
            if (ret != NVML_SUCCESS) {
                dev_err(data->misc_dev.this_device,  /* 修复：移除& */
                       "Failed to get NVLink state: %s\n", 
                       nvmlErrorString(ret));
                ret = -EIO;
                break;
            }
            info.is_active = (state == NVML_ENABLED) ? 1 : 0;
            
            /* 获取远程设备PCI信息 */
            if (info.is_active) {
                nvmlPciInfo_t pci;
                ret = nvmlDeviceGetNvLinkRemotePciInfo(device, info.link, &pci);
                if (ret == NVML_SUCCESS) {
                    strncpy(info.remote_bus_id, pci.busId, sizeof(info.remote_bus_id)-1);
                    info.remote_domain = pci.domain;
                    info.remote_bus = pci.bus;
                    info.remote_device = pci.device;
                    info.remote_function = pci.function;
                }
            }
            
            /* 将结果返回给用户空间 */
            if (copy_to_user((void __user *)arg, &info, sizeof(info))) {
                ret = -EFAULT;
                break;
            }
            ret = 0;
            break;
        }
        
        case TWCLML_IOCTL_GET_P2P_INFO: {
            twclml_p2p_info_t info;
            nvmlDevice_t device1, device2;
            nvmlGpuP2PStatus_t status;
            
            /* 从用户空间获取输入参数 */
            if (copy_from_user(&info, (void __user *)arg, sizeof(info))) {
                ret = -EFAULT;
                break;
            }
            
            /* 获取两个设备的句柄 */
            ret = nvmlDeviceGetHandleByIndex(info.device1_index, &device1);
            if (ret != NVML_SUCCESS) {
                dev_err(data->misc_dev.this_device,  /* 修复：移除& */
                       "Failed to get device1 handle: %s\n", 
                       nvmlErrorString(ret));
                ret = -EIO;
                break;
            }
            
            ret = nvmlDeviceGetHandleByIndex(info.device2_index, &device2);
            if (ret != NVML_SUCCESS) {
                dev_err(data->misc_dev.this_device,  /* 修复：移除& */
                       "Failed to get device2 handle: %s\n", 
                       nvmlErrorString(ret));
                ret = -EIO;
                break;
            }
            
            /* 获取P2P读能力 */
            ret = nvmlDeviceGetP2PStatus(device1, device2, 
                                       NVML_P2P_CAPS_INDEX_READ, &status);
            info.read_supported = (ret == NVML_SUCCESS && status == NVML_P2P_STATUS_OK) ? 1 : 0;
            
            /* 获取P2P写能力 */
            ret = nvmlDeviceGetP2PStatus(device1, device2, 
                                       NVML_P2P_CAPS_INDEX_WRITE, &status);
            info.write_supported = (ret == NVML_SUCCESS && status == NVML_P2P_STATUS_OK) ? 1 : 0;
            
            /* 获取NVLink上的P2P能力 */
            ret = nvmlDeviceGetP2PStatus(device1, device2, 
                                       NVML_P2P_CAPS_INDEX_NVLINK, &status);
            info.nvlink_supported = (ret == NVML_SUCCESS && status == NVML_P2P_STATUS_OK) ? 1 : 0;
            
            /* 获取P2P原子操作能力 */
            ret = nvmlDeviceGetP2PStatus(device1, device2, 
                                       NVML_P2P_CAPS_INDEX_ATOMICS, &status);
            info.atomics_supported = (ret == NVML_SUCCESS && status == NVML_P2P_STATUS_OK) ? 1 : 0;
            
            /* 将结果返回给用户空间 */
            if (copy_to_user((void __user *)arg, &info, sizeof(info))) {
                ret = -EFAULT;
                break;
            }
            ret = 0;
            break;
        }
        
        default:
            ret = -ENOTTY;
            break;
    }
    
out_unlock:
    mutex_unlock(&data->lock);
    return ret;
}

/* 文件操作结构体 */
static const struct file_operations twclml_fops = {
    .owner      = THIS_MODULE,
    .open       = twclml_open,
    .release    = twclml_release,
    .unlocked_ioctl = twclml_ioctl,
#ifdef CONFIG_COMPAT
    .compat_ioctl = twclml_ioctl,
#endif
};

/**
 * @brief platform驱动probe函数
 */
static int twclml_platform_probe(struct platform_device *pdev)
{
    int ret;
    
    dev_info(&pdev->dev, "TWCLML platform driver probing...\n");
    
    /* 分配驱动数据结构 */
    twclml_drv_data = devm_kzalloc(&pdev->dev, sizeof(*twclml_drv_data), GFP_KERNEL);
    if (!twclml_drv_data)
        return -ENOMEM;
    
    /* 初始化互斥锁 */
    mutex_init(&twclml_drv_data->lock);
    
    /* 初始化misc设备 */
    twclml_drv_data->misc_dev.name = TWCLML_DEVICE_NAME;
    twclml_drv_data->misc_dev.fops = &twclml_fops;
    twclml_drv_data->misc_dev.minor = MISC_DYNAMIC_MINOR;
    twclml_drv_data->initialized = 0;
    
    /* 注册misc设备 */
    ret = misc_register(&twclml_drv_data->misc_dev);
    if (ret) {
        dev_err(&pdev->dev, "Failed to register misc device: %d\n", ret);
        mutex_destroy(&twclml_drv_data->lock);
        return ret;
    }
    
    /* 将驱动数据与平台设备关联 */
    platform_set_drvdata(pdev, twclml_drv_data);
    
    dev_info(&pdev->dev, "TWCLML platform driver probed successfully\n");
    dev_info(&pdev->dev, "Device node created: %s\n", TWCLML_DEVICE_PATH);
    return 0;
}

/**
 * @brief platform驱动remove函数
 */
static int twclml_platform_remove(struct platform_device *pdev)
{
    struct twclml_driver_data *data = platform_get_drvdata(pdev);
    
    dev_info(&pdev->dev, "TWCLML platform driver removing...\n");
    
    /* 注销misc设备 */
    misc_deregister(&data->misc_dev);
    
    /* 销毁互斥锁 */
    mutex_destroy(&data->lock);
    
    dev_info(&pdev->dev, "TWCLML platform driver removed successfully\n");
    return 0;
}

/* platform设备匹配表 */
static const struct of_device_id twclml_of_match[] = {
    { .compatible = "twcl,twclml", },
    { /* sentinel */ }
};
MODULE_DEVICE_TABLE(of, twclml_of_match);

/* platform驱动结构体 */
static struct platform_driver twclml_platform_driver = {
    .probe  = twclml_platform_probe,
    .remove = twclml_platform_remove,
    .driver = {
        .name = "twclml",
        .of_match_table = twclml_of_match,
    },
};

/* 平台设备释放函数 */
static void twclml_platform_device_release(struct device *dev)
{
    pr_info("TWCLML platform device released\n");
}

/* 平台设备定义（添加release回调） */
static struct platform_device twclml_platform_device = {
    .name = "twclml",
    .id = -1,
    .dev = {
        .release = twclml_platform_device_release,
    },
};

/**
 * @brief 模块初始化函数
 */
static int __init twclml_driver_init(void)
{
    int ret;
    
    /* 注册platform设备 */
    ret = platform_device_register(&twclml_platform_device);
    if (ret) {
        pr_err("Failed to register platform device: %d\n", ret);
        return ret;
    }
    
    /* 注册platform驱动 */
    ret = platform_driver_register(&twclml_platform_driver);
    if (ret) {
        pr_err("Failed to register platform driver: %d\n", ret);
        platform_device_unregister(&twclml_platform_device);
        return ret;
    }
    
    pr_info("TWCLML driver initialized\n");
    return 0;
}

/**
 * @brief 模块退出函数
 */
static void __exit twclml_driver_exit(void)
{
    /* 注销platform驱动和设备 */
    platform_driver_unregister(&twclml_platform_driver);
    platform_device_unregister(&twclml_platform_device);
    
    pr_info("TWCLML driver exited\n");
}

module_init(twclml_driver_init);
module_exit(twclml_driver_exit);

MODULE_AUTHOR("Your Name");
MODULE_DESCRIPTION("TWCLML Platform-Misc Driver");
MODULE_LICENSE("GPL");
MODULE_ALIAS("platform:twclml");